import streamlit as st
from openai import OpenAI
from swarm import Agent, Swarm
import json
import os
import sys
import time
import datetime
# 获取当前文件所在目录的上一级目录的绝对路径
parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
# 将上一级目录添加到sys.path中
sys.path.append(parent_dir)
from back_end.Customs_agent_assistant.system_module.agents_instructions import triage_instructions, analysis_instructions
from back_end.Customs_agent_assistant.system_module.file_operations import load_json_file, add_record_to_json
from back_end.Customs_agent_assistant.system_module.qdrant_operations import query_docs,milvus_search
from back_end.Customs_agent_assistant.system_module.utilities import Evoke_hot_spot_analysis_capabilities
from back_end.Customs_agent_assistant.system_module.run_demo import pretty_format_messages,process_and_print_streaming_response,pretty_print_messages
import threading
from queue import Queue,Empty
input_queue = Queue()
output_queue = Queue()
# 初始化会话状态
def initialize_session_state():
    if "messages" not in st.session_state:
        st.session_state.messages = [{"role": "assistant", "content": "有什么需要帮忙的？"}]
    if "history" not in st.session_state:
        st.session_state.history = []

# 渲染历史对话记录侧边栏
def render_sidebar():
    with st.sidebar:
        st.markdown("### 历史对话记录")
        
        # 显示历史对话记录标签
        for i, history_record in enumerate(st.session_state.history):
            if st.button(history_record["tag"], key=f"history_{i}"):
                st.session_state.messages = history_record["messages"]
        
        # 添加新增会话的按钮
        if st.button("新增会话", use_container_width=True):
            if len(st.session_state.messages) > 1: 
                first_user_message = next((msg["content"] for msg in st.session_state.messages if msg["role"] == "user"), "新对话")
                st.session_state.history.append({
                    "tag": first_user_message[:20] + "..." if len(first_user_message) > 20 else first_user_message,
                    "messages": st.session_state.messages
                })
            st.session_state.messages = [{"role": "assistant", "content": "有什么需要帮忙的？"}]
            st.rerun()
        
        # 添加清空历史对话记录的按钮
        if st.button("清空历史对话记录", use_container_width=True):
            st.session_state.history = []
            st.rerun()

# 渲染默认图片框
def render_default_image():
    if len(st.session_state.messages) == 1 and st.session_state.messages[0]["role"] == "assistant":
        with st.container():
            col1, col2, col3 = st.columns([1, 3, 1])
            with col2:
                st.markdown("### 昨日热搜")
                with st.container(border=True):
                    st.image("image\image2.png", use_container_width=True)
                    if st.button("查看更多", use_container_width=True):
                        st.page_link(r"pages\display.py", label="查看更多")

# 渲染聊天记录
def render_chat_messages():
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])

# 处理用户输入
def handle_user_input():
    if prompt := st.chat_input("请输入您的问题："):
        st.session_state.messages.append({"role": "user", "content": prompt})
        with st.chat_message("user"):
            styled_message = get_chat_message_style(prompt, True)
            st.markdown(styled_message, unsafe_allow_html=True)
        input_queue.put(prompt)
        # LLM生成回复,替换assistant_response
        # xiugai
        #with st.chat_message("assistant"):
        #    assistant_response="NIHAO"
        #    st.markdown(assistant_response)

        #st.session_state.messages.append({"role": "assistant", "content": assistant_response})
                
# 渲染清空当前对话的按钮
def render_clear_current_chat_button():
    if st.session_state.messages:
        col1, col2 = st.columns([4, 1])
        with col2:
            if st.button("清空当前对话", use_container_width=True):
                st.session_state.messages = [{"role": "assistant", "content": "有什么需要帮忙的？"}]
                st.rerun()

api_key="sk-feebadbcb5654f5c8f9044c78f7c4548"
client = OpenAI(api_key=api_key,
                base_url="https://dashscope.aliyuncs.com/compatible-mode/v1")

file_path = 'D:/robotlearn/agent/Customs_agent_assistant/data/system_data/Keyword statistics.json'
context_variables = load_json_file(file_path)

# 定义一个函数用于将请求转移到分诊智能体
def transfer_to_triage():
    """当用户的请求需要转移到不同的智能体或不同的政策时，调用此函数。
    例如，当用户询问的内容不属于当前智能体处理范围时，调用此函数进行转移。
    """
    return triage_agent

# 定义一个函数用于将请求转移到报关记录查询智能体
def transfer_to_customs_record_query():
    """当用户的请求需要转移到报关记录查询智能体时，调用此函数。
    例如，当用户查询报关记录时，调用此函数进行转移。
    """
    return customs_record_query_agent


# 定义一个函数用于将请求转移到对话分析智能体
def transfer_to_Conversation_analytics():
    """当用户的请求已经确认结束后，需要转移到对话分析智能体来更新关键词信息，调用此函数。
    """
    return Conversation_analytics_agent

def transfer_to_hot_spot_analysis_agent():
    """当用户的请求需要转移到热点分析智能体时，调用此函数。
    """
    return Hot_spot_analysis_agent

def transfer_to_customs_knowledge_query_agent():
    """当用户的请求需要转移到海关知识查询智能体时，调用此函数。
    """
    return customs_knowledge_query_agent



triage_agent = Agent(
    name="Triage Agent",
    model="qwen-max-0919",  # 智能体名称：分诊智能体
    instructions=triage_instructions,  # 调用分诊指令，根据上下文帮助处理
    functions=[transfer_to_customs_record_query,
    transfer_to_hot_spot_analysis_agent,
    transfer_to_customs_knowledge_query_agent,
    ],  # 定义可调用的函数，分别转移到航班修改和行李丢失
)

customs_record_query_agent = Agent(
    name='Customs_Record_Query',
    model="qwen-max-0919",
    instructions="""
        你是中国海关12360热线的一名报关记录查询智能体。
        1.请询问用户想要查询的海关报关记录的关键词，并调用query_docs进行查询。
        2.如果用户的请求需要转移到不同的智能体或不同的政策时，调用transfer_to_triage函数。
        3.如果客户没有进一步问题，调用 'transfer_to_Conversation_analytics' 函数。
    """,
    functions=[query_docs,transfer_to_triage,transfer_to_Conversation_analytics],
)

customs_knowledge_query_agent = Agent(
    name='Customs_Knowledge_Query',
    model="qwen-max-0919",
    instructions="""
        你是中国海关12360热线的一名报关知识查询智能体。
        1.请询问用户想要查询的海关报关知识的问题，并作为query输入到milvus_search函数进行查询。
        2.如果用户的请求需要转移到不同的智能体或不同的政策时，调用transfer_to_triage函数。
        3.如果客户没有进一步问题，调用 'transfer_to_Conversation_analytics' 函数。
    """,
    functions=[milvus_search,transfer_to_triage,transfer_to_Conversation_analytics],
)

Conversation_analytics_agent = Agent(
    name="Conversation analytics",
    model = "qwen-max-0919",
    instructions=analysis_instructions,
    functions=[add_record_to_json,transfer_to_triage]
)

Hot_spot_analysis_agent = Agent(
    name="Hot spot analysis",
    model="qwen-max-0919",
    instructions="""
        你是中国海关12360热线的一名热点分析智能体。
        1.请询问用户是否确定要进行热点分析，如果得到用户的肯定回复，就调用Evoke_hot_spot_analysis_capabilities函数唤起功能。
        2.如果客户没有进一步问题，调用 'transfer_to_Conversation_analytics' 函数
        3.如果用户的请求需要转移到不同的智能体或不同的政策时，调用transfer_to_triage函数。
        注意：不要进行多余的输出，你只需要负责判断是否调用该函数。
    """,
    functions=[Evoke_hot_spot_analysis_capabilities,transfer_to_Conversation_analytics,transfer_to_triage]
)

def run_demo_loop_to_backend1(
        openai_client,
        starting_agent,
        context_variables=None,
        input_queue=None,
        socketio_instance=None,
        stream=False,
        debug=False,
         # 添加一个参数用于接收输入队列
) -> None:
    client = Swarm(openai_client)
    print("Starting Swarm CLI 🐝")
    print('Type "exit" or "quit" to leave the chat.')

    messages = []
    agent = starting_agent

    if input_queue is None:
        raise ValueError("input_queue parameter cannot be None")

    while True:
        try:
            user_input = input_queue.get(block=True, timeout=1)  # 从队列中获取输入内容，设置超时时间
            user_input = user_input.strip()
            if user_input.lower() in {"exit", "quit"}:
                with open('knowledgeBase/Keyword statistics.json', 'w', encoding='UTF-8') as f:
                    json.dump(context_variables, f, ensure_ascii=False)
                print("Exiting chat. Goodbye!")
                break
            messages.append({"role": "user", "content": user_input})

            response = client.run(
                agent=agent,
                messages=messages,
                context_variables=context_variables or {},
                stream=stream,
                debug=debug,
                max_turns=100
            )

            if stream:
                response = process_and_print_streaming_response(response)
            else:
                pretty_print_messages(response.messages)
                formatted_messages = pretty_format_messages(response.messages)
                output_queue.put(formatted_messages)
                #for formatted_message in formatted_messages:
                 #   socketio_instance.emit('new_message', {'message': formatted_message})

            messages.extend(response.messages)
            agent = response.agent
            context_variables = response.context_variables
        except Empty:
            continue

def get_chat_message_style(message, is_user):
    timestamp = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    base_style = 'padding: 10px; border-radius: 5px; margin-bottom: 10px;'
    if is_user:
        bg_color = "#f0f0f0"
        icon = "👤"
    else:
        bg_color = "#e0e0e0"
        icon = "🤖"
    style = f'<div style="background-color: {bg_color}; {base_style}"><span style="margin-right: 5px;">{icon}</span><b>{ "用户" if is_user else "助手" } ({timestamp})</b>: {message}</div>'
    # 添加响应式样式，示例中当屏幕宽度小于600px时调整字体大小和消息宽度
    style += '<style>@media (max-width: 600px) { div { font-size: 14px; width: 90%; }}</style>'
    return style

# 主函数
def main():
    st.title("你好，我是ICA")
    initialize_session_state()
    render_sidebar()
    render_default_image()
    render_chat_messages()
    
    render_clear_current_chat_button()
    handle_user_input()



if __name__ == "__main__":
    st.title("你好，我是ICA")
    initialize_session_state()
    render_sidebar()
    render_default_image()
    render_chat_messages()
    render_clear_current_chat_button()
    handle_user_input()
    loop_thread = threading.Thread(target=run_demo_loop_to_backend1, args=(client, triage_agent, context_variables, input_queue))
    loop_thread.start()
    while True:
        try:
            # 尝试从输出队列获取数据
            assistant_response = output_queue.get(block=False)
            with st.chat_message("assistant"):
                styled_message = get_chat_message_style(assistant_response, False)
                st.markdown(styled_message, unsafe_allow_html=True)

            st.session_state.messages.append({"role": "assistant", "content": assistant_response})
        except Empty:
            time.sleep(0.1)
            continue
    #main()